rm(list = ls())
library(tidyverse)
library(broom)
library(haven)
library(polycor)
library(wCorr)
library(quanteda)
library(quanteda.textplots)
require(glmnet)
require(quanteda.textstats)
library(ggcorrplot)
library(tidymodels)

load("../working/survey.Rdata")


## Treatment effects on outcomes

# Wave 1 respondents

get_position_effects <- function(my_data = just_w1, w = "weight", outcome_wave = "w2", subset = NULL) {
  
  if(!is.null(subset)) my_data <- my_data[subset,]
  
  if(is.null(w)) {
    
    my_data$tmp_weight <- 1
    w <- "tmp_weight"
    
  }
  
  if(outcome_wave == "w2"){
    
    my_data$zero_hours <-  my_data$zero_hours_w2
    my_data$high_tax <-  my_data$high_tax_w2
    my_data$minimum_wage <-  my_data$minimum_wage_w2
    my_data$trans_rights <-  my_data$trans_rights_w2
    my_data$offensive_speech <-  my_data$offensive_speech_w2
    my_data$unemployment_support <-  my_data$unemployment_support_w2
  
  }
  
  zero_coefs <- broom::tidy(lm(scale(zero_hours) ~ treat, my_data, weights = get(w)), conf.int = TRUE)
  high_tax_coefs <- broom::tidy(lm(scale(high_tax) ~ treat, my_data, weights = get(w)), conf.int = TRUE)
  minimum_wage_coefs <- broom::tidy(lm(scale(minimum_wage) ~ treat, my_data, weights = get(w)), conf.int = TRUE)
  trans_rights_coefs <- broom::tidy(lm(scale(trans_rights) ~ treat, my_data, weights = get(w)), conf.int = TRUE)
  offensive_speech_coefs <- broom::tidy(lm(scale(offensive_speech) ~ treat, my_data, weights = get(w)), conf.int = TRUE)
  unemployment_support_coefs <- broom::tidy(lm(scale(unemployment_support) ~ treat, my_data, weights = get(w)), conf.int = TRUE)
  
  mean_est <- my_data %>%
    select(c(Id, treat, all_of(w), zero_hours, high_tax, minimum_wage, unemployment_support, trans_rights, offensive_speech)) %>%
    pivot_longer(-c(treat, all_of(w), Id)) %>%
    filter(!is.na(value)) %>%
    estimatr::lm_robust(scale(value) ~ treat + name, data = ., weights = get(w), clusters = Id) %>% tidy()
  
  treat_effects <- bind_rows(
    tibble(var = "Zero Hours", zero_coefs[2,c("estimate", "conf.low", "conf.high")]),
    tibble(var = "High Tax", high_tax_coefs[2,c("estimate", "conf.low", "conf.high")]),
    tibble(var = "Minimum Wage", minimum_wage_coefs[2,c("estimate", "conf.low", "conf.high")]),
    tibble(var = "Transgender Rights", trans_rights_coefs[2,c("estimate", "conf.low", "conf.high")]),
    tibble(var = "Offensive Speech", offensive_speech_coefs[2,c("estimate", "conf.low", "conf.high")]),
    tibble(var = "Unemployment Support", unemployment_support_coefs[2,c("estimate", "conf.low", "conf.high")]),
    tibble(var = "Mean effect", mean_est[2,c("estimate", "conf.low", "conf.high")])
  )
  
  treat_effects
  
}

effects_wave_1 <- get_position_effects(my_data = just_w1, 
                                       w = "weight", outcome_wave = "w1")

effects_wave_2_panel <- get_position_effects(my_data = just_panel, 
                                             w = "weight", outcome_wave = "w2")

effects_wave_2_topup <- get_position_effects(my_data = just_topup, 
                                             w = "weight", outcome_wave = "w2")

effects_wave_1$wave <- "Sample One, Wave One"
effects_wave_2_panel$wave <- "Sample One, Wave Two"
effects_wave_2_topup$wave <- "Sample Two"

average_treat_effects <- rbind(effects_wave_1[7,],
                               effects_wave_2_panel[7,],
                               effects_wave_2_topup[7,])

average_treat_effects$var <- "Zero Hours"

average_treat_effects <- average_treat_effects %>% 
  mutate(wave = factor(wave, levels = c("Sample One, Wave One", "Sample One, Wave Two", "Sample Two")))

bind_rows(effects_wave_1, effects_wave_2_panel, effects_wave_2_topup) %>%
  filter(var != "Mean effect") %>%
  mutate(var = factor(var, levels = effects_wave_1$var[order(effects_wave_1$estimate)]),
         var = droplevels(var),
         wave = factor(wave, levels = c("Sample One, Wave One", "Sample One, Wave Two", "Sample Two"))) %>%
  ggplot(aes(x = estimate, xmin = conf.low, xmax = conf.high, y = var, col = wave, fill = wave)) + 
  geom_pointrange(position = position_dodge(width = .2)) + 
  facet_wrap(~wave) + 
  geom_rect(data = average_treat_effects, ymax = 100, ymin = -100, alpha = .2, lwd = 0.0001) + 
  geom_vline(aes(xintercept = estimate, col = wave), data = average_treat_effects) + 
  theme_bw() + 
  xlab("Effect of reason giving on attitudes") + 
  ylab("") + 
  theme(panel.background = element_rect(fill = "transparent"),
        plot.background = element_rect(fill = "transparent", color = NA),
        legend.background = element_rect(fill = "transparent", color = NA),
        legend.box.background = element_rect(fill = "transparent", color = NA),
        legend.key = element_blank(),
        legend.position = "bottom") +
  geom_vline(xintercept = 0, linetype = 2) + 
  scale_color_manual("", values = c("black", "black", "black")) + 
  scale_fill_manual("", values = c("black", "black", "black")) + 
  guides(fill = "none", color = "none") 
ggsave("../out/outcomes/left_right.pdf", width = 10, height = 5)
